from glob import glob
from os.path import join as pathJoin
from PIL import Image
import numpy as np
from tqdm import tqdm

def findDatasetPaths(labels, path='./datasets/'):
    ''' 查找文件集路径 '''
    datasetPaths = []
    for label in labels:
        datasetPaths.append(glob(pathJoin(path, str(label), '*')))
    return datasetPaths

def readImage(path, format_size=(64, 64), color_mode='RGB'):
    ''' 读取图片像素信息 '''
    img = Image.open(path).resize(format_size, resample=Image.BOX).convert(color_mode)
    return img

def readDatasets(labels, path='./datasets/', format_size=(64, 64)):
    ''' 读取数据集 '''
    datasetPaths = findDatasetPaths(labels, path)
    datasets = []
    for i in range(len(datasetPaths)):
        datasets.append([])
        # for path in datasetPaths[i]:
        for path in tqdm(datasetPaths[i], desc='Class: {}'.format(i), ncols=80):
            img = readImage(path, format_size)
            datasets[i].append(np.array(img, dtype='uint8'))
    return datasets

def getDatas(datasets, dtype='uint8'):
    ''' 从训练集中得到输入数据 '''
    datas = []
    for dataset in datasets:
        for data in dataset:
            datas.append(data)
    return np.array(datas, dtype=dtype)

def getTargets(datasets):
    ''' 从训练集中得到输出数据 '''
    targets = []
    for i, size in enumerate(map(len, datasets)):
        targets.extend([i] * size)
    return np.array(targets, dtype='int32')
